#include <vector>
#include <random>
#include <ctime>
#include <algorithm>

using namespace std;

/*
528. 按权重随机选择
已解答
中等
相关标签
相关企业
给你一个 下标从 0 开始 的正整数数组 w ，其中 w[i] 代表第 i 个下标的权重。

请你实现一个函数 pickIndex ，它可以 随机地 从范围 [0, w.length - 1] 内（含 0 和 w.length - 1）选出并返回一个下标。选取下标 i 的 概率 为 w[i] / sum(w) 。

例如，对于 w = [1, 3]，挑选下标 0 的概率为 1 / (1 + 3) = 0.25 （即，25%），而选取下标 1 的概率为 3 / (1 + 3) = 0.75（即，75%）。
 

示例 1：

输入：
["Solution","pickIndex"]
[[[1]],[]]
输出：
[null,0]
解释：
Solution solution = new Solution([1]);
solution.pickIndex(); // 返回 0，因为数组中只有一个元素，所以唯一的选择是返回下标 0。
示例 2：

输入：
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
输出：
[null,1,1,1,1,0]
解释：
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // 返回 1，返回下标 1，返回该下标概率为 3/4 。
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 0，返回下标 0，返回该下标概率为 1/4 。

由于这是一个随机问题，允许多个答案，因此下列输出都可以被认为是正确的:
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
......
诸若此类。
 

提示：

1 <= w.length <= 104
1 <= w[i] <= 105
pickIndex 将被调用不超过 104 次
*/

// 法一
class Solution1 {
public:
    // 处理一下累积和数组 
    Solution1(vector<int>& w) {
        preFixSum.push_back(w[0]);
        for(int i = 1; i < w.size(); ++i) {
            preFixSum.push_back(preFixSum[i - 1] + w[i]);
        }
        total = preFixSum.back();   // 总权重的和
        std::srand(std::time(0));   
    }
    // 随机选择下标
    int pickIndex() {
        // 区间随机数
        int tar = std::rand() % total + 1;
        // 二分查找区间
        int left = 0;
        int right = preFixSum.size() - 1;
        while(left < right) {
            int mid = left + (right - left) / 2;
            if(preFixSum[mid] < tar)    left = mid + 1;
            else    right = mid;
        }
        return left;
    }
private:
    vector<int> preFixSum;
    int total;
};

/**
 * Your Solution object will be instantiated and called as such:
 * Solution* obj = new Solution(w);
 * int param_1 = obj->pickIndex();
 */
    
// 法二
class Solution2 {
public:
    Solution2(vector<int>& w) {
        int m = w.size();
        sum.resize(m + 1, 0);
        for (int i = 1; i <= m; i++) {
            sum[i] = sum[i - 1] + w[i - 1];
        }
    }

    int pickIndex() {
        uniform_int_distribution<int> uni(0, sum.back() - 1);  // sum.back() 直接取得最后一个元素
        int num = uni(rng);
        // upper_bound 返回的是大于 num 的最小值的下标
        int a = upper_bound(sum.begin(), sum.end(), num) - sum.begin();
        return a - 1;  // 返回区间内的实际索引
    }

private:
    vector<int> sum; // 存储累积和
    mt19937 rng{random_device{}()}; // 随机数生成器
};
